package mybatis.plus;

import com.baomidou.mybatisplus.generator.config.ConstVal;
import com.baomidou.mybatisplus.generator.config.builder.ConfigBuilder;
import com.baomidou.mybatisplus.generator.config.po.TableInfo;
import com.baomidou.mybatisplus.generator.engine.AbstractTemplateEngine;
import com.baomidou.mybatisplus.toolkit.StringUtils;
import org.apache.velocity.Template;
import org.apache.velocity.VelocityContext;
import org.apache.velocity.app.Velocity;
import org.apache.velocity.app.VelocityEngine;

import java.io.BufferedWriter;
import java.io.FileOutputStream;
import java.io.OutputStreamWriter;
import java.util.Map;
import java.util.Properties;

/**
 * @author Arnold Yand
 * @time 2018/3/27 14:36
 */
public class UserDefineTemplateEngine extends AbstractTemplateEngine {
  private static final String DOT_VM = ".vm";
  private VelocityEngine velocityEngine;

  @Override
  public UserDefineTemplateEngine init(ConfigBuilder configBuilder) {
    super.init(configBuilder);
    if (null == velocityEngine) {
      Properties p = new Properties();
      p.setProperty(ConstVal.VM_LOADPATH_KEY, ConstVal.VM_LOADPATH_VALUE);
      p.setProperty(Velocity.FILE_RESOURCE_LOADER_PATH, "");
      p.setProperty(Velocity.ENCODING_DEFAULT, ConstVal.UTF8);
      p.setProperty(Velocity.INPUT_ENCODING, ConstVal.UTF8);
      p.setProperty("file.resource.loader.unicode", "true");
      velocityEngine = new VelocityEngine(p);
    }
    return this;
  }
  @Override
  public void writer(Map<String, Object> objectMap, String templatePath, String outputFile) throws Exception {
    if (StringUtils.isEmpty(templatePath)) {
      return;
    }
    StringBuffer serviceName = new StringBuffer(((TableInfo)objectMap.get("table")).getServiceName());
    serviceName.replace(0, 1, "");
    serviceName.setCharAt(0, Character.toLowerCase(serviceName.charAt(0)));
    objectMap.put("serviceAnnotationValue", serviceName.toString());
    objectMap.put("versionFieldName", "version");
    Template template = velocityEngine.getTemplate(templatePath, ConstVal.UTF8);
    FileOutputStream fos = new FileOutputStream(outputFile);
    BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(fos, ConstVal.UTF8));
    template.merge(new VelocityContext(objectMap), writer);
    writer.close();
    logger.debug("模板:" + templatePath + ";  文件:" + outputFile);
  }

  @Override
  public String templateFilePath(String filePath) {
    if (null == filePath || filePath.contains(DOT_VM)) {
      return filePath;
    }
    StringBuilder fp = new StringBuilder();
    fp.append(filePath).append(DOT_VM);
    return fp.toString();
  }
}
